#include "cpu/flash_attn_v1.h"
#include "cpu/naive_attn.h"
#include "perf.h"

#include "gtest/gtest.h"
#include <vector>
#include <random>

using namespace std;

void RandomInit(vector<float>& v, int size) {
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<> dis(-1.0, 1.0);
    for (int i = 0; i < size; ++i) {
        v[i] = dis(gen);
    }
}


TEST(FlashAttnV1CpuTest, Test) {
    using DATA_TYPE = float;

    constexpr int kH = 2;
    constexpr int kN = 3096;
    constexpr int kD = 128;
    constexpr int kBr = 127;
    constexpr int kBc = 127;

    vector<DATA_TYPE> q(kN * kD);
    RandomInit(q, kN * kD);
    
    vector<DATA_TYPE> k(kN * kD);
    RandomInit(k, kN * kD);

    vector<DATA_TYPE> v(kN * kD);
    RandomInit(v, kN * kD);
    
    vector<DATA_TYPE> o_naive(kN * kD, 0.0);
    vector<DATA_TYPE> o_v1(kN * kD, 0.0);

    {
        PERF(naive_attn);
        naive_attn<DATA_TYPE, kN, kD>(q, k, v, o_naive);
    }


    {
        PERF(flash_attn_v1);
        flash_attn_v1<DATA_TYPE, kH, kN, kD, kBr, kBc>(q, k, v, o_v1);
    }

    for (int i = 0; i < kN * kD; ++i) {
        if(std::isnan(o_naive[i]) or \
           std::isnan(o_v1[i]) or \
           (std::abs(o_naive[i] - o_v1[i]) > 10e-5))
        {
            printf("Error: %f %f, row: %d, col: %d\n", o_naive[i], o_v1[i], i / kD,
                i % kD);
        }
    }
}